// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/sparksql/size.h>

#include <pollux/core/query_config.h>

#include <pollux/functions/macros.h>
#include <pollux/functions/registerer.h>

namespace kumo::pollux::functions::sparksql {
    namespace {
        template<typename TExecParams>
        struct Size {
            POLLUX_DEFINE_FUNCTION_TYPES(TExecParams);

            template<typename TInput>
            MELON_ALWAYS_INLINE void initialize(
                const std::vector<TypePtr> & /*inputTypes*/,
                const core::QueryConfig & /*config*/,
                const TInput * /*input*/,
                const bool *legacySizeOfNull) {
                if (legacySizeOfNull == nullptr) {
                    POLLUX_USER_FAIL("Constant legacySizeOfNull is expected.");
                }
                legacySizeOfNull_ = *legacySizeOfNull;
            }

            template<typename TInput>
            MELON_ALWAYS_INLINE bool callNullable(
                int32_t &out,
                const TInput *input,
                const bool * /*legacySizeOfNull*/) {
                if (input == nullptr) {
                    if (legacySizeOfNull_) {
                        out = -1;
                        return true;
                    }
                    return false;
                }
                out = input->size();
                return true;
            }

        private:
            // If true, returns -1 for null input. Otherwise, returns null.
            bool legacySizeOfNull_;
        };
    } // namespace

    void registerSize(const std::string &prefix) {
        register_function<Size, int32_t, Array<Any>, bool>({prefix});
        register_function<Size, int32_t, Map<Any, Any>, bool>({prefix});
    }
} // namespace kumo::pollux::functions::sparksql
